import os
import base64
import numpy as np
import torch
import torchaudio
from fastapi import FastAPI, Request, HTTPException, Body, Form, File, UploadFile, Header, Depends
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, HTMLResponse
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any, Union
from io import BytesIO
import re
import time
import json
import uvicorn
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
from enum import Enum

# 设置模型 - 使用本地离线模型
model_path = "models/iic/SenseVoiceSmall"
vad_model_path = "models/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch"

# 打印模型路径信息
print(f"=== 模型加载信息 ===")
print(f"模型路径: {os.path.abspath(model_path)}")
print(f"VAD模型路径: {os.path.abspath(vad_model_path)}")
print(f"模型路径是否存在: {os.path.exists(model_path)}")
print(f"VAD模型路径是否存在: {os.path.exists(vad_model_path)}")
if os.path.exists(model_path):
    print(f"模型目录内容: {os.listdir(model_path)}")
if os.path.exists(vad_model_path):
    print(f"VAD模型目录内容: {os.listdir(vad_model_path)}")
print(f"=== 模型加载信息结束 ===")

# 加载模型
print("开始加载模型...")
model = AutoModel(model=model_path,
                 vad_model=vad_model_path,
                 vad_kwargs={"max_single_segment_time": 30000},
                 trust_remote_code=False)
print("模型加载完成!")

# 定义语言枚举类
class Language(str, Enum):
    auto = "auto"
    zh = "zh"
    en = "en"
    yue = "yue"
    ja = "ja"
    ko = "ko"
    nospeech = "nospeech"

regex = r"<\|.*\|>"

# 定义情感和事件标记字典，用于后处理
emo_dict = {
    "<|HAPPY|>": "😊",
    "<|SAD|>": "😔",
    "<|ANGRY|>": "😡",
    "<|NEUTRAL|>": "",
    "<|FEARFUL|>": "😰",
    "<|DISGUSTED|>": "🤢",
    "<|SURPRISED|>": "😮",
}

event_dict = {
    "<|BGM|>": "🎼",
    "<|Speech|>": "",
    "<|Applause|>": "👏",
    "<|Laughter|>": "😀",
    "<|Cry|>": "😭",
    "<|Sneeze|>": "🤧",
    "<|Breath|>": "",
    "<|Cough|>": "🤧",
}

emoji_dict = {
    "<|nospeech|><|Event_UNK|>": "❓",
    "<|zh|>": "",
    "<|en|>": "",
    "<|yue|>": "",
    "<|ja|>": "",
    "<|ko|>": "",
    "<|nospeech|>": "",
    "<|HAPPY|>": "😊",
    "<|SAD|>": "😔",
    "<|ANGRY|>": "😡",
    "<|NEUTRAL|>": "",
    "<|BGM|>": "🎼",
    "<|Speech|>": "",
    "<|Applause|>": "👏",
    "<|Laughter|>": "😀",
    "<|FEARFUL|>": "😰",
    "<|DISGUSTED|>": "🤢",
    "<|SURPRISED|>": "😮",
    "<|Cry|>": "😭",
    "<|EMO_UNKNOWN|>": "",
    "<|Sneeze|>": "🤧",
    "<|Breath|>": "",
    "<|Cough|>": "😷",
    "<|Sing|>": "",
    "<|Speech_Noise|>": "",
    "<|withitn|>": "",
    "<|woitn|>": "",
    "<|GBG|>": "",
    "<|Event_UNK|>": "",
}

lang_dict = {
    "<|zh|>": "<|lang|>",
    "<|en|>": "<|lang|>",
    "<|yue|>": "<|lang|>",
    "<|ja|>": "<|lang|>",
    "<|ko|>": "<|lang|>",
    "<|nospeech|>": "<|lang|>",
}

emo_set = {"😊", "😔", "😡", "😰", "🤢", "😮"}
event_set = {"🎼", "👏", "😀", "😭", "🤧", "😷",}

# 文本格式化函数
def format_str(s):
    for sptk in emoji_dict:
        s = s.replace(sptk, emoji_dict[sptk])
    return s

def format_str_v2(s):
    sptk_dict = {}
    for sptk in emoji_dict:
        sptk_dict[sptk] = s.count(sptk)
        s = s.replace(sptk, "")
    emo = "<|NEUTRAL|>"
    for e in emo_dict:
        if sptk_dict[e] > sptk_dict[emo]:
            emo = e
    for e in event_dict:
        if sptk_dict[e] > 0:
            s = event_dict[e] + s
    s = s + emo_dict[emo]

    for emoji in emo_set.union(event_set):
        s = s.replace(" " + emoji, emoji)
        s = s.replace(emoji + " ", emoji)
    return s.strip()

def format_str_v3(s):
    def get_emo(s):
        return s[-1] if s[-1] in emo_set else None
    def get_event(s):
        return s[0] if s[0] in event_set else None

    s = s.replace("<|nospeech|><|Event_UNK|>", "❓")
    for lang in lang_dict:
        s = s.replace(lang, "<|lang|>")
    s_list = [format_str_v2(s_i).strip(" ") for s_i in s.split("<|lang|>")]
    new_s = " " + s_list[0]
    cur_ent_event = get_event(new_s)
    for i in range(1, len(s_list)):
        if len(s_list[i]) == 0:
            continue
        if get_event(s_list[i]) == cur_ent_event and get_event(s_list[i]) != None:
            s_list[i] = s_list[i][1:]
        #else:
        cur_ent_event = get_event(s_list[i])
        if get_emo(s_list[i]) != None and get_emo(s_list[i]) == get_emo(new_s):
            new_s = new_s[:-1]
        new_s += s_list[i].strip().lstrip()
    new_s = new_s.replace("The.", " ")
    return new_s.strip()

# 创建OpenAI兼容的API应用
app = FastAPI(title="SenseVoice OpenAI API Compatible Server")

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 简单的API密钥验证
API_KEYS = {"sk-sensevoice": "default"}

def verify_api_key(authorization: str = Header(None)):
    if not authorization:
        return None
    
    api_key = authorization.replace("Bearer ", "")
    if api_key not in API_KEYS and api_key != "":
        raise HTTPException(status_code=401, detail="无效的API密钥")
    
    return api_key

# OpenAI API 兼容模型
class AudioTranscriptionRequest(BaseModel):
    model: str = Field(default="sensevoice-small")
    file: str  # base64编码的音频文件
    language: Optional[str] = Field(default="auto")
    response_format: Optional[str] = Field(default="text")
    temperature: Optional[float] = Field(default=0)
    prompt: Optional[str] = Field(default="")

class AudioTranscriptionResponse(BaseModel):
    text: str

class Message(BaseModel):
    role: str
    content: Union[str, List[Dict[str, Any]]]

class ChatCompletionRequest(BaseModel):
    model: str
    messages: List[Message]
    temperature: Optional[float] = 0
    top_p: Optional[float] = 1
    n: Optional[int] = 1
    stream: Optional[bool] = False
    max_tokens: Optional[int] = None
    presence_penalty: Optional[float] = 0
    frequency_penalty: Optional[float] = 0
    language: Optional[str] = "auto"

class Choice(BaseModel):
    index: int
    message: Message
    finish_reason: str = "stop"

class Usage(BaseModel):
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int

class ChatCompletionResponse(BaseModel):
    id: str
    object: str = "chat.completion"
    created: int
    model: str
    choices: List[Choice]
    usage: Usage

class ModelList(BaseModel):
    object: str = "list"
    data: List[Dict[str, Any]]

class CompletionRequest(BaseModel):
    model: str
    prompt: Union[str, List[str]]
    max_tokens: Optional[int] = 16
    temperature: Optional[float] = 1.0
    language: Optional[str] = "auto"
    audio: Optional[str] = None  # Base64编码的音频数据

class CompletionResponse(BaseModel):
    id: str
    object: str
    created: int
    model: str
    choices: List[Dict[str, Any]]
    usage: Dict[str, int]

def process_audio(audio_data, language="auto"):
    """处理音频数据并返回转录结果"""
    try:
        # 确保 audio_data 是字节类型
        if not isinstance(audio_data, bytes):
            print(f"警告: audio_data 不是字节类型，而是 {type(audio_data)}，尝试转换")
            audio_data = bytes(audio_data)
        
        print(f"处理音频数据: {len(audio_data)} 字节, 语言: {language}")
        
        # 加载音频
        try:
            data_or_path_or_list, audio_fs = torchaudio.load(BytesIO(audio_data))
            print(f"音频加载成功: 采样率 {audio_fs} Hz, 形状 {data_or_path_or_list.shape}")
        except Exception as e:
            print(f"音频加载失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
        
        # 预处理音频
        data_or_path_or_list = data_or_path_or_list.mean(0).numpy()
        data_or_path_or_list = data_or_path_or_list.astype(np.float32)
        if len(data_or_path_or_list.shape) > 1:
            data_or_path_or_list = data_or_path_or_list.mean(-1)
            print(f"音频降维为单声道: 形状 {data_or_path_or_list.shape}")
        
        # 重采样到 16kHz（如果需要）
        if audio_fs != 16000:
            print(f"重采样音频: {audio_fs} Hz -> 16000 Hz")
            try:
                resampler = torchaudio.transforms.Resample(audio_fs, 16000)
                input_wav_t = torch.from_numpy(data_or_path_or_list).to(torch.float32)
                data_or_path_or_list = resampler(input_wav_t[None, :])[0, :].numpy()
                print(f"重采样后音频形状: {data_or_path_or_list.shape}")
            except Exception as e:
                print(f"重采样失败: {str(e)}")
                import traceback
                traceback.print_exc()
                return None
        
        # 执行推理
        merge_vad = True
        print(f"开始推理: 语言={language}, merge_vad={merge_vad}")
        try:
            result = model.generate(
                input=data_or_path_or_list,
                cache={},
                language=language,
                use_itn=True,
                batch_size_s=60,
                merge_vad=merge_vad
            )
            print(f"推理完成: 结果长度 {len(result) if result else 0}")
        except Exception as e:
            print(f"推理失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
        
        if not result or len(result) == 0:
            print("推理结果为空")
            return None
        
        # 处理结果
        text = result[0]["text"]
        print(f"原始文本: {text[:100]}{'...' if len(text) > 100 else ''}")
        
        try:
            formatted_text = format_str_v3(text)
            print(f"格式化文本: {formatted_text[:100]}{'...' if len(formatted_text) > 100 else ''}")
        except Exception as e:
            print(f"格式化文本失败: {str(e)}")
            formatted_text = text
        
        try:
            clean_text = re.sub(regex, "", text, 0, re.MULTILINE)
            print(f"清理后文本: {clean_text[:100]}{'...' if len(clean_text) > 100 else ''}")
        except Exception as e:
            print(f"清理文本失败: {str(e)}")
            clean_text = text
        
        return {
            "raw_text": text,
            "formatted_text": formatted_text,
            "clean_text": clean_text
        }
        
    except Exception as e:
        print(f"处理音频时出错: {str(e)}")
        import traceback
        traceback.print_exc()
        return None

@app.get("/v1/models", response_model=ModelList)
async def list_models(api_key: str = Depends(verify_api_key)):
    """列出可用的模型"""
    return {
        "object": "list",
        "data": [
            {
                "id": "sensevoice-small",
                "object": "model",
                "created": int(time.time()),
                "owned_by": "organization-owner",
            }
        ]
    }

@app.post("/v1/audio/transcriptions")
async def create_transcription(request: Request, api_key: str = Depends(verify_api_key)):
    """处理音频转录请求，使用原始请求处理以避免解码错误"""
    try:
        audio_data = None
        model = "sensevoice-small"
        language = "auto"
        response_format = "text"
        
        # 检查内容类型
        content_type = request.headers.get("content-type", "")
        
        if "application/json" in content_type:
            # 处理 JSON 请求
            body = await request.body()
            try:
                data = json.loads(body)
                model = data.get("model", "sensevoice-small")
                file_data = data.get("file", "")
                language = data.get("language", "auto")
                response_format = data.get("response_format", "text")
                
                # 解码 base64 音频
                try:
                    # 如果 base64 字符串包含前缀，去掉前缀
                    if "," in file_data:
                        file_data = file_data.split(",", 1)[1]
                    audio_data = base64.b64decode(file_data)
                except Exception as e:
                    print(f"Base64 解码失败: {str(e)}")
                    return JSONResponse(
                        status_code=400,
                        content={"error": f"无效的 base64 编码音频数据: {str(e)}"}
                    )
            except json.JSONDecodeError as e:
                return JSONResponse(
                    status_code=400,
                    content={"error": f"无效的 JSON 数据: {str(e)}"}
                )
        elif "multipart/form-data" in content_type:
            # 处理表单数据
            try:
                form = await request.form()
                model = form.get("model", "sensevoice-small")
                language = form.get("language", "auto")
                response_format = form.get("response_format", "text")
                
                # 获取文件
                file = form.get("file")
                if file is None:
                    return JSONResponse(
                        status_code=400,
                        content={"error": "未提供音频文件"}
                    )
                
                if isinstance(file, UploadFile):
                    # 如果是上传的文件对象，读取其内容
                    audio_data = await file.read()
                    print(f"读取到上传文件: {file.filename}, 大小: {len(audio_data)} 字节")
                elif isinstance(file, str):
                    # 如果是 base64 字符串
                    try:
                        # 如果 base64 字符串包含前缀，去掉前缀
                        if "," in file:
                            file = file.split(",", 1)[1]
                        audio_data = base64.b64decode(file)
                    except Exception as e:
                        print(f"Base64 解码失败: {str(e)}")
                        return JSONResponse(
                            status_code=400,
                            content={"error": f"无效的 base64 编码音频数据: {str(e)}"}
                        )
                else:
                    # 如果是其他类型，打印类型信息并尝试转换为字节
                    print(f"接收到未知类型的文件数据: {type(file)}")
                    try:
                        # 尝试读取内容（如果是类文件对象）
                        if hasattr(file, 'read') and callable(file.read):
                            audio_data = await file.read() if callable(getattr(file, 'read', None)) else file.read()
                        elif hasattr(file, 'file') and hasattr(file.file, 'read'):
                            # 可能是某种文件包装器
                            audio_data = await file.file.read() if callable(getattr(file.file, 'read', None)) else file.file.read()
                        else:
                            # 尝试直接使用
                            audio_data = bytes(file)
                    except Exception as e:
                        print(f"读取文件数据失败: {str(e)}")
                        traceback.print_exc()
                        return JSONResponse(
                            status_code=400,
                            content={"error": f"无法读取音频数据: {str(e)}"}
                        )
            except Exception as e:
                print(f"处理表单数据失败: {str(e)}")
                import traceback
                traceback.print_exc()
                return JSONResponse(
                    status_code=400,
                    content={"error": f"处理表单数据失败: {str(e)}"}
                )
        else:
            # 处理原始二进制数据
            try:
                body = await request.body()
                # 检查是否有查询参数
                params = request.query_params
                model = params.get("model", "sensevoice-small")
                language = params.get("language", "auto")
                response_format = params.get("response_format", "text")
                
                # 假设主体是音频数据
                audio_data = body
            except Exception as e:
                print(f"读取请求体失败: {str(e)}")
                return JSONResponse(
                    status_code=400,
                    content={"error": f"读取请求体失败: {str(e)}"}
                )
        
        # 检查是否获取到音频数据
        if not audio_data:
            return JSONResponse(
                status_code=400,
                content={"error": "未提供有效的音频数据"}
            )
        
        # 处理音频
        result = process_audio(audio_data, language)
        
        if not result:
            return JSONResponse(
                content={"text": ""},
                status_code=200
            )
        
        # 根据请求的响应格式返回结果
        if response_format == "text":
            return JSONResponse(
                content={"text": result["formatted_text"]},
                status_code=200
            )
        else:
            return JSONResponse(
                content={"text": result["formatted_text"]},
                status_code=200
            )
            
    except Exception as e:
        print(f"处理请求时出错: {str(e)}")
        import traceback
        traceback.print_exc()
        return JSONResponse(
            status_code=400,
            content={"error": f"处理音频时出错: {str(e)}"}
        )

@app.post("/v1/chat/completions")
async def chat_completions(request: Request, api_key: str = Depends(verify_api_key)):
    """处理聊天完成请求，支持音频处理"""
    try:
        # 获取请求体
        body = await request.body()
        
        # 解析JSON
        try:
            data = json.loads(body)
        except json.JSONDecodeError as e:
            return JSONResponse(
                status_code=400,
                content={"error": f"无效的JSON数据: {str(e)}"}
            )
        
        # 获取必要的字段
        model = data.get("model", "sensevoice-small")
        messages = data.get("messages", [])
        language = data.get("language", "auto")
        
        if not messages:
            return JSONResponse(
                status_code=400,
                content={"error": "未提供消息"}
            )
        
        # 获取最后一条消息
        last_message = messages[-1]
        if last_message.get("role") != "user":
            return JSONResponse(
                status_code=400,
                content={"error": "最后一条消息必须来自用户"}
            )
        
        content = last_message.get("content", "")
        
        # 处理多模态内容
        audio_data = None
        if isinstance(content, list):
            for content_part in content:
                if content_part.get("type") == "audio" and "audio_data" in content_part:
                    try:
                        audio_data_str = content_part["audio_data"]
                        # 如果 base64 字符串包含前缀，去掉前缀
                        if "," in audio_data_str:
                            audio_data_str = audio_data_str.split(",", 1)[1]
                        audio_data = base64.b64decode(audio_data_str)
                        print(f"从多模态内容中提取音频数据: {len(audio_data)} 字节")
                        break
                    except Exception as e:
                        print(f"Base64解码失败: {str(e)}")
                        import traceback
                        traceback.print_exc()
                        return JSONResponse(
                            status_code=400,
                            content={"error": f"无效的base64编码音频数据: {str(e)}"}
                        )
        else:
            # 尝试从内容中提取base64音频
            audio_pattern = r"base64,([\w+/=]+)"
            match = re.search(audio_pattern, content)
            
            if match:
                try:
                    audio_data = base64.b64decode(match.group(1))
                    print(f"从文本内容中提取音频数据: {len(audio_data)} 字节")
                except Exception as e:
                    print(f"Base64解码失败: {str(e)}")
                    import traceback
                    traceback.print_exc()
                    return JSONResponse(
                        status_code=400,
                        content={"error": f"无效的base64编码音频数据: {str(e)}"}
                    )
        
        if not audio_data:
            return JSONResponse(
                status_code=200,
                content={
                    "id": f"chatcmpl-{int(time.time())}",
                    "object": "chat.completion",
                    "created": int(time.time()),
                    "model": model,
                    "choices": [
                        {
                            "index": 0,
                            "message": {
                                "role": "assistant",
                                "content": "请提供有效的音频数据进行转录。"
                            },
                            "finish_reason": "stop"
                        }
                    ],
                    "usage": {
                        "prompt_tokens": len(str(content)),
                        "completion_tokens": 0,
                        "total_tokens": len(str(content))
                    }
                }
            )
        
        # 处理音频
        result = process_audio(audio_data, language)
        
        if not result:
            transcription_text = "无法识别音频内容"
        else:
            # 构建响应文本
            transcription_text = f"音频转录结果：\n\n{result['clean_text']}\n\n富文本结果（包含情感和事件标记）：\n\n{result['formatted_text']}"
        
        # 构建响应
        response_content = transcription_text
        response_length = len(response_content)
        
        return JSONResponse(
            status_code=200,
            content={
                "id": f"chatcmpl-{int(time.time())}",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": model,
                "choices": [
                    {
                        "index": 0,
                        "message": {
                            "role": "assistant",
                            "content": response_content
                        },
                        "finish_reason": "stop"
                    }
                ],
                "usage": {
                    "prompt_tokens": len(str(content)),
                    "completion_tokens": response_length,
                    "total_tokens": len(str(content)) + response_length
                }
            }
        )
        
    except Exception as e:
        print(f"处理请求时出错: {str(e)}")
        import traceback
        traceback.print_exc()
        return JSONResponse(
            status_code=400,
            content={"error": f"处理请求时出错: {str(e)}"}
        )

@app.post("/v1/completions", response_model=CompletionResponse)
async def create_completion(request: CompletionRequest, api_key: str = Depends(verify_api_key)):
    """处理文本补全请求（用于音频转文本）"""
    # 检查是否提供了音频数据
    if not request.audio:
        raise HTTPException(status_code=400, detail="需要提供音频数据")
    
    try:
        # 解码Base64音频数据
        audio_data_str = request.audio
        # 如果 base64 字符串包含前缀，去掉前缀
        if "," in audio_data_str:
            audio_data_str = audio_data_str.split(",", 1)[1]
        audio_data = base64.b64decode(audio_data_str)
        
        # 处理音频
        result = process_audio(audio_data, request.language)
        
        # 提取文本结果
        text = ""
        if result:
            text = result["clean_text"]
        
        # 构建OpenAI格式的响应
        return {
            "id": f"cmpl-{int(time.time())}",
            "object": "text_completion",
            "created": int(time.time()),
            "model": request.model,
            "choices": [
                {
                    "text": text,
                    "index": 0,
                    "logprobs": None,
                    "finish_reason": "stop"
                }
            ],
            "usage": {
                "prompt_tokens": 0,
                "completion_tokens": len(text.split()) if text else 0,
                "total_tokens": len(text.split()) if text else 0
            }
        }
    except Exception as e:
        print(f"处理音频时出错: {str(e)}")
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"处理音频时出错: {str(e)}")

# 主页
@app.get("/", response_class=HTMLResponse)
async def root():
    return """
    <!DOCTYPE html>
    <html>
        <head>
            <meta charset=utf-8>
            <title>SenseVoice OpenAI API</title>
        </head>
        <body>
            <h1>SenseVoice OpenAI API</h1>
            <p>这是SenseVoice的OpenAI API兼容接口。</p>
            <p>可用的接口：</p>
            <ul>
                <li><a href="/docs">/docs</a> - API文档</li>
                <li>/v1/models - 列出可用模型</li>
                <li>/v1/completions - 文本补全（音频转文本）</li>
                <li>/v1/chat/completions - 聊天补全（音频转文本）</li>
                <li>/v1/audio/transcriptions - 音频文件转录</li>
            </ul>
        </body>
    </html>
    """

if __name__ == "__main__":
    port = int(os.getenv("PORT", "7000"))
    print(f"启动SenseVoice OpenAI API兼容服务，监听端口: {port}")
    uvicorn.run(app, host="0.0.0.0", port=port) 